{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Supervised Fine-Tuning (SFT) at scale with DeepSpeed\n",
    "\n",
    "This guide provides a step-by-step workflow for supervised fine-tuning the [`Qwen/Qwen2.5-32B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-32B-Instruct) model on a multi-GPU Anyscale cluster. You use LLaMA-Factory for the training framework and `DeepSpeed` to efficiently manage memory and scale the training process.\n",
    "\n",
    "SFT is a technique to adapt a pre-trained model to specific tasks. By showing the model high-quality examples of instructions and their desired outputs, you teach it to follow new instructions more accurately.\n",
    "\n",
    "## Step 1: Set up your environment\n",
    "\n",
    "### Dependencies\n",
    "First, ensure your environment has the correct libraries. Start with a pre-built container image and install LLaMA-Factory and DeepSpeed on top of it.\n",
    "\n",
    "Recommended container image:\n",
    "```bash\n",
    "anyscale/ray-llm:2.48.0-py311-cu128\n",
    "```\n",
    "\n",
    "Execute the following commands to install the required packages and optional tools for experiment tracking and faster model downloads:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Install the specific version of LLaMA-Factory\n",
    "pip install -q llamafactory==0.9.3\n",
    "\n",
    "# Install DeepSpeed for large-scale training\n",
    "pip install -q deepspeed==0.16.9\n",
    "\n",
    "# (Optional) For experiment tracking with Weights & Biases\n",
    "pip install -q wandb==0.21.3\n",
    "\n",
    "# (Optional) For accelerated model downloads from Hugging Face\n",
    "pip install -q hf_transfer==0.1.9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model and compute resources\n",
    "\n",
    "DeepSpeed ZeRO-3 partitions parameters, gradients, and optimizer states across multiple GPUs, enabling supervised fine-tuning (SFT) of 30B+ LLMs on just 4 GPUs.\n",
    "\n",
    "| Item | Value |\n",
    "|------|-------|\n",
    "| **Base model** | [`Qwen/Qwen2.5-32B-Instruct`](https://huggingface.co/Qwen/Qwen2.5-32B-Instruct) |\n",
    "| **Worker Nodes** | 4 × L40S / 4 x A100-40G |\n",
    "\n",
    "## Step 2: Prepare the dataset\n",
    "\n",
    "### Understand the dataset\n",
    "This tutorial uses [`glaive_toolcall_en_demo`](https://huggingface.co/datasets/zuol/glaive_toolcall_en_demo/tree/main), a dataset designed to teach models how to use tools (also known as function calling).\n",
    "\n",
    "This dataset contains conversational examples where the model needs to interact with external tools. Each entry includes:\n",
    "* `conversations`: A turn-by-turn log between a human and the gpt assistant.\n",
    "* `tools`: A JSON schema describing the functions the model can call.\n",
    "\n",
    "**Note**: The `conversations` may include special turns like function_call (the model deciding to call a tool) and observation (the result returned from the tool). This structure is ideal for teaching the model sophisticated tool-use behavior. To maintain role alignment in ShareGPT format, you must follow a strict turn order: `human` and `observation` (tool output) must appear in odd-numbered positions, while `gpt` and `function_call` must appear in even-numbered positions. The model learns to generate the content in the `gpt` and `function_call` turns.\n",
    "\n",
    "**Dataset example**\n",
    "\n",
    "```json\n",
    "{\n",
    "\"conversations\": [\n",
    "    {\n",
    "    \"from\": \"human\",\n",
    "    \"value\": \"Hi, I have some ingredients and I want to cook something. Can you help me find a recipe?\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"gpt\",\n",
    "    \"value\": \"Of course! I can help you with that. Please tell me what ingredients you have.\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"human\",\n",
    "    \"value\": \"I have chicken, bell peppers, and rice.\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"function_call\",\n",
    "    \"value\": \"{\\\"name\\\": \\\"search_recipes\\\", \\\"arguments\\\": {\\\"ingredients\\\": [\\\"chicken\\\", \\\"bell peppers\\\", \\\"rice\\\"]}}\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"observation\",\n",
    "    \"value\": \"{\\\"recipes\\\": [{\\\"name\\\": \\\"Chicken and Bell Pepper Stir Fry\\\", \\\"ingredients\\\": [\\\"chicken\\\", \\\"bell peppers\\\", \\\"rice\\\"], \\\"instructions\\\": \\\"Cut the chicken into small pieces. Slice the bell peppers. Cook the rice. Stir fry the chicken and bell peppers. Serve over rice.\\\"}, {\\\"name\\\": \\\"Chicken and Rice Casserole\\\", \\\"ingredients\\\": [\\\"chicken\\\", \\\"bell peppers\\\", \\\"rice\\\"], \\\"instructions\\\": \\\"Cook the chicken and rice separately. Mix them together with the bell peppers in a casserole dish. Bake until golden brown.\\\"}]}\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"gpt\",\n",
    "    \"value\": \"I found two recipes for you. The first one is \\\"Chicken and Bell Pepper Stir Fry\\\". The instructions are: Cut the chicken into small pieces. Slice the bell peppers. Cook the rice. Stir fry the chicken and bell peppers. Serve over rice. The second one is \\\"Chicken and Rice Casserole\\\". The instructions are: Cook the chicken and rice separately. Mix them together with the bell peppers in a casserole dish. Bake until golden brown. Which one would you like to try?\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"human\",\n",
    "    \"value\": \"They both sound delicious, but I think I'll try the stir fry. Can you order the ingredients for me?\"\n",
    "    },\n",
    "    {\n",
    "    \"from\": \"gpt\",\n",
    "    \"value\": \"I'm sorry, but as an AI, I don't have the capability to perform external tasks such as ordering ingredients. However, I can help you find more recipes or provide cooking tips if you need.\"\n",
    "    }\n",
    "],\n",
    "\"tools\": \"[{\\\"name\\\": \\\"search_recipes\\\", \\\"description\\\": \\\"Search for recipes based on ingredients\\\", \\\"parameters\\\": {\\\"type\\\": \\\"object\\\", \\\"properties\\\": {\\\"ingredients\\\": {\\\"type\\\": \\\"array\\\", \\\"items\\\": {\\\"type\\\": \\\"string\\\"}, \\\"description\\\": \\\"The ingredients to search for\\\"}}, \\\"required\\\": [\\\"ingredients\\\"]}}]\"\n",
    "}\n",
    "```\n",
    "\n",
    "### Register the dataset\n",
    "\n",
    "To specify new datasets that are accessible across Ray worker nodes, you must first add a **`dataset_info.json`** to **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)** such as `/mnt/cluster_storage`. This configuration file acts as a central registry for all your datasets. It maps a custom name to your dataset file location, format, and column structure. \n",
    "\n",
    "If you plan to run SFT fine-tuning on the `glaive_toolcall_en_demo` dataset, first complete the setup steps below. Ensure that you place the dataset files in a storage location that all workers can access (for example, a shared mount or object storage). Avoid storing large files on the head node.\n",
    "\n",
    "`dataset_info.json`\n",
    "```json\n",
    "{\n",
    "  \"my_glaive_toolcall_en_demo\": {\n",
    "      \"file_name\": \"/mnt/cluster_storage/glaive_toolcall_en_demo.json\",\n",
    "      \"formatting\": \"sharegpt\",\n",
    "      \"columns\": {\n",
    "          \"messages\": \"conversations\",\n",
    "          \"tools\": \"tools\"\n",
    "      }\n",
    "  }\n",
    "}\n",
    "```\n",
    "\n",
    "For a more detailed dataset preparation and formatting guide, see [Choose your data format](https://docs.anyscale.com/llm/fine-tuning/data-preparation#sft)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Make sure all files are accessible to worker nodes\n",
    "# Create a copy of the data in /mnt/cluster_storage\n",
    "wget https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/datasets/sharegpt/glaive_toolcall_en_demo.json -O /mnt/cluster_storage/glaive_toolcall_en_demo.json\n",
    "# Create a copy of the dataset registry in /mnt/cluster_storage\n",
    "cp ../dataset-configs/dataset_info.json /mnt/cluster_storage/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create the fine-tuning config (SFT with DeepSpeed)\n",
    "\n",
    "Next, create the main YAML configuration file—the master recipe for your fine-tuning job. It specifies the base model, the fine-tuning method (LoRA), the dataset, training hyperparameters, cluster resources, and more.\n",
    "\n",
    "**Important notes:**\n",
    "- **W&B tracking:** To track experiments with Weights & Biases (W&B), set `report_to: wandb` in the config and provide `WANDB_API_KEY` in the runtime environment. If you don't want to use W&B, set `report_to: none` to avoid errors.\n",
    "- **Access and paths:** The YAML only needs to be on the **head node**, but any referenced paths (`dataset_dir`, `output_dir`) must reside on storage **reachable by all workers** (for example, `/mnt/cluster_storage/`).\n",
    "- **Gated models:** If your base model has gated access (for example, Llama) on Hugging Face, set `HF_TOKEN` in the runtime environment.\n",
    "- **GPU selection and placement:** The config uses a 4xL40S node (`ananyscale/accelerator_shape:4xL40S`) so that all 4 GPUs are on the same machine, which is important for efficient DeepSpeed ZeRO-3 communication. You can switch to other multi-GPU nodes such as `4xA100-40GB` or any other node type with comparable or more VRAM, depending on your cloud availability.\n",
    "\n",
    "### Configure LLaMA-Factory with Ray\n",
    "\n",
    "**Note**: To customize the training configuration, edit `train-configs/sft_lora_deepspeed.yaml`. \n",
    "\n",
    "```yaml\n",
    "# sft_lora_deepspeed.yaml\n",
    "\n",
    "### model\n",
    "model_name_or_path: Qwen/Qwen2.5-32B-Instruct\n",
    "trust_remote_code: true\n",
    "\n",
    "### method\n",
    "stage: sft\n",
    "do_train: true\n",
    "finetuning_type: lora\n",
    "lora_rank: 8\n",
    "lora_target: all\n",
    "\n",
    "### deepspeed\n",
    "deepspeed: /mnt/cluster_storage/ds_z3_config.json # path to the DeepSpeed config\n",
    "\n",
    "### dataset\n",
    "dataset: my_glaive_toolcall_en_demo\n",
    "dataset_dir: /mnt/cluster_storage\n",
    "\n",
    "template: qwen\n",
    "cutoff_len: 1024\n",
    "max_samples: 1000\n",
    "overwrite_cache: true\n",
    "preprocessing_num_workers: 16\n",
    "\n",
    "### output\n",
    "output_dir: qwen2.5_32b_lora_sft\n",
    "logging_steps: 5\n",
    "save_steps: 50\n",
    "plot_loss: true\n",
    "report_to: wandb # or none\n",
    "\n",
    "### train\n",
    "per_device_train_batch_size: 1 # Adjust this depending on your GPU memory and sequence length\n",
    "gradient_accumulation_steps: 4\n",
    "num_train_epochs: 3.0\n",
    "learning_rate: 1.0e-4\n",
    "bf16: true\n",
    "lr_scheduler_type: cosine\n",
    "warmup_ratio: 0.1\n",
    "ddp_timeout: 180000000\n",
    "\n",
    "### ray\n",
    "ray_run_name: qwen2.5_32b_lora_sft\n",
    "ray_storage_path: /mnt/cluster_storage/\n",
    "ray_num_workers: 4  # Number of GPUs to use\n",
    "resources_per_worker:\n",
    "  GPU: 1\n",
    "  # accelerator_type:L40S: 0.001            # Use this to simply specify a GPU type (may place GPUs on separate nodes).\n",
    "  anyscale/accelerator_shape:4xL40S: 0.001  # Prefer this for DeepSpeed so all 4 GPUs are on the same node.\n",
    "  # See https://docs.ray.io/en/master/ray-core/accelerator-types.html#accelerator-types for a full list of accelerator types.\n",
    "ray_init_kwargs:\n",
    "  runtime_env:\n",
    "    env_vars:\n",
    "      # If using wandb for experiments tracking\n",
    "      WANDB_API_KEY: <your_wandb_token>\n",
    "      # If using gated models like meta-llama/Llama-3.1-8B-Instruct\n",
    "      # HF_TOKEN: <your_huggingface_token>\n",
    "      # If hf_transfer is installed\n",
    "      HF_HUB_ENABLE_HF_TRANSFER: '1'\n",
    "```\n",
    "\n",
    "**Note:**\n",
    "This configuration assumes `4xL40S` GPUs are available in your cloud environment. If not, you can substitute with `4xA100-40G` (or another supported accelerator with similar VRAM).\n",
    "\n",
    "### DeepSpeed configuration\n",
    "DeepSpeed is an open-source deep-learning optimization library developed by Microsoft, aimed at enabling large-model training. Higher ZeRO stages (1→3) and enabling CPU offload reduce GPU VRAM usage, but might cause slower training.\n",
    "\n",
    "To enable DeepSpeed, create a separate Deepspeed config in the **[storage shared across nodes](https://docs.anyscale.com/configuration/storage#shared)**. and reference it from your main training yaml config with:\n",
    "\n",
    "```yaml\n",
    "deepspeed: /mnt/cluster_storage/ds_z3_config.json\n",
    "```\n",
    "\n",
    "Below is a sample ZeRO-3 config:\n",
    "\n",
    "`ds_z3_config.json`\n",
    "```json\n",
    "{\n",
    "\"train_batch_size\": \"auto\",\n",
    "\"train_micro_batch_size_per_gpu\": \"auto\",\n",
    "\"gradient_accumulation_steps\": \"auto\",\n",
    "\"gradient_clipping\": \"auto\",\n",
    "\"zero_allow_untested_optimizer\": true,\n",
    "\"fp16\": {\n",
    "    \"enabled\": \"auto\",\n",
    "    \"loss_scale\": 0,\n",
    "    \"loss_scale_window\": 1000,\n",
    "    \"initial_scale_power\": 16,\n",
    "    \"hysteresis\": 2,\n",
    "    \"min_loss_scale\": 1\n",
    "},\n",
    "\"bf16\": {\n",
    "    \"enabled\": \"auto\"\n",
    "},\n",
    "\"zero_optimization\": {\n",
    "    \"stage\": 3,\n",
    "    \"overlap_comm\": false,\n",
    "    \"contiguous_gradients\": true,\n",
    "    \"sub_group_size\": 1e9,\n",
    "    \"reduce_bucket_size\": \"auto\",\n",
    "    \"stage3_prefetch_bucket_size\": \"auto\",\n",
    "    \"stage3_param_persistence_threshold\": \"auto\",\n",
    "    \"stage3_max_live_parameters\": 1e9,\n",
    "    \"stage3_max_reuse_distance\": 1e9,\n",
    "    \"stage3_gather_16bit_weights_on_model_save\": true\n",
    "}\n",
    "}\n",
    "```\n",
    "\n",
    "For a more detailed guide on acceleration and optimization methods including DeepSpeed on Ray, see [Speed and memory optimizations](https://docs.anyscale.com/llm/fine-tuning/speed-and-memory-optimizations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Create a copy of the DeepSpeed configuration file in /mnt/cluster_storage\n",
    "cp ../deepspeed-configs/ds_z3_config.json /mnt/cluster_storage/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Train and monitor\n",
    "\n",
    "**Note**: If you installed Weights & Biases, set `WANDB_API_KEY` in the runtime environment. Otherwise, set `report_to: none` in `sft_lora_deepspeed.yaml` to avoid `api_token not set` errors.\n",
    "\n",
    "With all configurations in place, you can launch fine-tuning or post-training in one of two ways:\n",
    "\n",
    "### Option A: Run from a workspace (quick start)\n",
    "\n",
    "The `USE_RAY=1` prefix tells LLaMA-Factory to run in distributed mode on the Ray cluster attached to your workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "USE_RAY=1 llamafactory-cli train ../train-configs/sft_lora_deepspeed.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option B: Run as an Anyscale job (production)\n",
    "\n",
    "For longer or production runs, submit the training as an **Anyscale job**. Jobs run outside your interactive session for better stability, retries, and durable logs. You package LLaMA-Factory and other libraries in a container image and launch with a short job config. See [Run LLaMA-Factory as an Anyscale job](https://docs.anyscale.com/llm/fine-tuning/llamafactory-jobs) for the step-by-step guide.\n",
    "\n",
    "### Monitoring with Weights & Biases (WandB)\n",
    "If you enabled Weights & Biases (with `report_to: wandb` in the training config YAML file), you can monitor your training job in real-time. Look for the training loss to decrease steadily, which indicates the model is learning.\n",
    "\n",
    "**Weights & Biases example**\n",
    "\n",
    "![WandB](https://anyscale-public-materials.s3.us-west-2.amazonaws.com/llm-finetuning/llama-factory/3.2.1/3.2.1-wandb.png)\n",
    "\n",
    "For a more detailed guide on tracking experiments with other tools such as Weights & Biases or MLflow, see [Observability and tracking](https://docs.anyscale.com/llm/fine-tuning/observability-and-tracking).\n",
    "\n",
    "## Step 5: Locate checkpoints\n",
    "\n",
    "Ray Train writes checkpoints under `ray_storage_path/ray_run_name`. In this example run, the path is: `/mnt/cluster_storage/qwen2.5_32b_lora_sft`. \n",
    "\n",
    "Inside, you see a **trainer session** directory named like:\n",
    "`TorchTrainer_8c6a5_00000_0_2025-09-09_09-53-45/`.\n",
    "\n",
    "- Ray Train creates `TorchTrainer_*` **when the trainer starts**; the suffix encodes a short run ID and the **start timestamp**.\n",
    "- Within that directory, Ray Train names checkpoints `checkpoint_000xxx/`, where the number is the saved ordered checkpoints.\n",
    "\n",
    "Control the save cadence with `save_strategy` and `save_steps`. For instructions on how to resume interrupted training with `resume_from_checkpoint` and more, see [Understand the artifacts directory](https://docs.anyscale.com/llm/fine-tuning/checkpointing#artifacts-directory).\n",
    "\n",
    "## Step 6: Export the model\n",
    "\n",
    "If you use LoRA, you can keep the base model and adapters separate for [multi-LoRA deployment](https://docs.anyscale.com/llm/serving/multi-lora) or [merge the adapters](https://docs.anyscale.com/llm/fine-tuning/checkpointing#merge-lora) into the base model for low-latency inference. \n",
    "\n",
    "For full fine-tuning or freeze-tuning, export the fine-tuned model directly.\n",
    "\n",
    "You may optionally apply [post-training quantization](https://docs.anyscale.com/llm/fine-tuning/checkpointing#ptq) on merged or full models before serving."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
